#include "VitterHuffmanCodec.h"
#include <bitset>
#include <iostream>


using namespace Huffman;

VitterHuffmanCodec::VitterHuffmanCodec(std::string inputFile, std::string outputFile) 
	: HuffmanCodec(inputFile, outputFile)
{
	
	// intialize the symbol tree
	symbolTree = new Tree();
}

VitterHuffmanCodec::~VitterHuffmanCodec(void)
{
	// empty destructor
}

void VitterHuffmanCodec::encode() 
{
	
	while(!inputDataStream -> isEof()) 
	{		
		encodeSymbol();
	}


	closeStreams();
}

void VitterHuffmanCodec::decode()
{
	int eocFound = 0;
	while(!eocFound) 
	{		
		eocFound = decodeSymbol();
	}
		
	closeStreams();
}

void VitterHuffmanCodec::closeStreams() {

	inputDataStream -> close();	
	outputDataStream -> flush();
	outputDataStream -> close();
}


void VitterHuffmanCodec::encodeSymbol() 
{
	
	Node* node;
	int encodedSymbol;
	char symbolBytes;
	char symbol;

	// read entire byte from input
	symbol = inputDataStream -> read();

	// finds the coresponding node in the symbol tree or the NYT node
	node = symbolTree -> findNode(symbol);	
	/*encodedSymbol = getSymbolCodeFromNode(node);*/

	
	//TODO back to this logic, eliminate EncodedSymbol class
	/* prevents from writing the first ESC code - siply output the code instea
	if(encodedSymbol > 0) 
	{
		wrtieSymbolBytes(encodedSymbol);
	} */

	//gdzies jest bald w kodzie, na koncu pliku wejsciowego do kodera zaczytuje jeszcze jeden bajt (11111111) dopiero wychodzi
	if ( inputDataStream -> isEof() ){ 
		writeEncodedSymbol(getEncodedSymbolCodeFromNode(Tree::EOC));
		return;
	}

	if(!node) // == Tree::NYT
	{						
		writeEncodedSymbol(getEncodedSymbolCodeFromNode(Tree::NYT)); /*writeSymbol(getSymbolCodeFromNode(Tree::NYT));*/		
		writeSymbol(symbol);
	} else {		
		writeEncodedSymbol(getEncodedSymbolCodeFromNode(node)); //getSymbolCodeFromNode(node)
	}

	symbolTree -> updateTreeVitter(symbol);		
}

int VitterHuffmanCodec::getSymbolCodeFromNode(Node* node) {

	int result = 0;	

	/* Ten if chyba jest zbedny ale nie mam teraz pewnosci */
	if ( node -> getParent() == 0 && node == symbolTree -> getRoot() ){
		return node->getSymbol(); //NYT
	}

	while(node -> getParent() != 0) {		
		result <<= 1;
		result |= node -> getParity();

		node = node -> getParent();		
	}
	
	return result;
}

EncodedSymbol* VitterHuffmanCodec::getEncodedSymbolCodeFromNode(Node* node) {

	int result = 0;	
	int bitCount = 0;

	if ( node -> getParent() == 0 ){ //&& node == symbolTree -> getRoot() 
		return new EncodedSymbol(node -> getSymbol(), 1); //NYT
	}

	while(node -> getParent() != 0) {		
		result <<= 1;
		result |= node -> getParity();

		node = node -> getParent();		
		bitCount++;
	}
	
	return new EncodedSymbol(result, bitCount);
}

char VitterHuffmanCodec::getSymbolBytes(int symbol) {

	// TODO not sure what exactly this should do, but in my understanding 
	// this method should just return the symbol from input, eventually 
	// in a form of byte array
	// TODO consider using std::bitset instead of int values
	return symbol;
}

void VitterHuffmanCodec::writeSymbolBytes(int bytes) {

	while(bytes > 0) {
		outputDataStream -> write(bytes & 1);

		bytes >>= 1;
	}
}

void VitterHuffmanCodec::writeEncodedSymbol(EncodedSymbol* symbol) {
	
	outputDataStream -> writeEncodedSymbol(symbol);

}

void VitterHuffmanCodec::writeSymbol(char symbol) {
	
	outputDataStream -> write(symbol);

}

int VitterHuffmanCodec::decodeSymbol() 
{	
	bool val;
	Node* node = symbolTree -> getRoot();
	
	int decodedSymbol;

	// iterates over the bits of the value
	for(int i = 7; i >= 0; i--) 
	{

		// reads a single bit from the input
		val = inputDataStream -> readBit();				

		if(val) 
		{
			node = node -> getRight();
		} else 
		{
			node = node -> getLeft();
		}

		// checks if the current , if was at begining of function but that was wrong
		if(node -> isLeaf()) 
		{
			break;
		}		
		
	}

	if(node == symbolTree -> EOC) {  //nastepnego bajtu juz nie bedzie, a to co dalej w aktualnym bajcie to same 0		
		return 1;
	}

	if(node == symbolTree -> NYT) 
	{
		if(!outputDataStream -> isEof()) {
			decodedSymbol = readSymbolBytes();			
		}
	} else {
		decodedSymbol = node -> getSymbol();
	}

	
	writeSymbol(decodedSymbol);
	symbolTree -> updateTreeVitter(decodedSymbol);
	return 0;
}

char VitterHuffmanCodec::readSymbolBytes() 
{
	return inputDataStream -> read();
}